import time
from Enet2 import Enet, UData,YData
from datetime import datetime
import sys


# 导入自定义函数和CANFD SDK动态链接库>>>
from os.path import dirname, join
import sys
sys.path.append(join(dirname(__file__), "src"))
sys.path.append(dirname(__file__))
from src import *


from imu import IMU
import time
imu=IMU()
time.sleep(2)




# Initialize Enet
enet = Enet()
enet.ReceiveY()
now=datetime.now()




loopPeriod=control_period
cnt=0

time.sleep(1)


site_force = np.zeros(6)
filtered_site_force = np.zeros(12)
force_bias = np.zeros(12)
filtered_imu_angle = np.zeros(3)


l = 0.3 # 头的宽度
roT_h = np.array([[0, 1, 0, 0],[0, 0, 1, 0],[1, 0, 0, -l],[0, 0, 0, 1]]) #头中心点到右臂o点的变换阵
loT_h = np.array([[0, 1, 0, 0],[0, 0, -1, 0],[-1, 0, 0, -l],[0, 0, 0, 1]]) #头中心点到左臂o点的变换阵
hT_e = np.eye(4)

box_quat = np.array([1,0,0,0])
imu.angle = [0,0,0]
clip_v = 0.1
clip_force_0 = 10
clip_k = 1e-5
getbox_flag = 0
force_sensor_cali_flag = 0


maxwidth = 0.6
minwidth = 0.05
t0 = 4
t_cali = 2
t1 = (maxwidth - minwidth)/clip_v + 2 # 夹取所需时间
t2 = 12 #抬起箱子所需时间
t3 = 20 #放下箱子所需时间

box_pos0 = np.array([0, 0.7, -0.68])
box_pos_lift = np.array([0, 0.5, -0.3])
box_pos_robo_end = np.array([0,0.05, -0.78])
box_width_robo_end = 0.8


try:
    
    loopStartTime=time.time()
    while True:
        
        cnt+=1
        start_time = time.time()
        t=start_time-loopStartTime

        # 2.1 ReadY from YPool
        start_time = time.time()

        y_data = enet.ReadY()
        
        q_real_degree = np.array(y_data.pos)
        q_real_degree_l = q_real_degree[0:6]
        q_real_degree_r = q_real_degree[6:12]
        q_real_int_l = trans_q_degree2intq(q_real_degree_l)
        q_real_int_r = trans_q_degree2intq(q_real_degree_r)

        q_real_float_l = trans_intq2q(q_real_int_l)
        q_real_float_r = trans_intq2q(q_real_int_r)

        

        # 校准力传感器
        while (time.time() - loopStartTime)<(t0+t_cali) and (time.time()-loopStartTime)>t0:
            if force_sensor_cali_flag == 0:
                force_bias = np.array(y_data.F)
                force_sensor_cali_flag = 1
            time.sleep(0.01)
            force_bias = (np.array(y_data.F) + force_bias.copy())/2


        site_force_raw = np.array(y_data.F) - force_bias #减去校准量
        filtered_site_force = alpha * site_force_raw + (1 - alpha) * filtered_site_force
        site_force = filtered_site_force.copy()

        # 双手挤压力
        clip_force = -(site_force[2] + site_force[8]) 
        

        if t<t0:
            # 先到预定位置
            box_pos = box_pos0
            box_width = maxwidth
            box_width_0 = maxwidth

        elif t<(t0 + t_cali + t1):

            # 开始夹取
            if clip_force <= clip_force_0 and getbox_flag == 0:
                box_width_0 -= loopPeriod*clip_v
            else:
                getbox_flag = 1

            if getbox_flag == 1:
                box_width_0 = box_width_0.copy() + clip_k*(clip_force - clip_force_0)
            
            box_width_0 = np.clip(box_width_0, minwidth, maxwidth)
            box_width = box_width_0

            
        elif t<(t0 + t_cali + t1 + t2 + t3) and getbox_flag == 1:
            
            box_width =  box_width_0 + 20*clip_k*(clip_force - 2.5*clip_force_0)
            box_width = np.clip(box_width, minwidth, box_width_0)

            if t <(t0 + t_cali + t1 + t2):
                box_pos = box_pos0 + np.clip(2*(t - t0 - t_cali - t1)/t2, 0, 1)*(box_pos_lift - box_pos0)
            else:
                box_pos = box_pos_lift + np.clip(2*(t - t0 - t_cali - t1 - t2)/t3, 0, 1) * (box_pos0 - box_pos_lift)


        else:
            if t < (t0 + t_cali + t1 + t2 + t3 + 1):
                box_pos = box_pos0
                box_width = maxwidth
            else:
                box_pos = box_pos_robo_end
                box_width = box_width_robo_end



        if len(imu.angle) == 3:
            imu_angle_raw = np.array([imu.angle[0],imu.angle[1],imu.angle[2]])
        filtered_imu_angle = alpha * imu_angle_raw + (1 - alpha) * filtered_imu_angle
        anglebias = np.pi*np.clip(filtered_imu_angle.copy()[0], -30, 30)/180

        head_pos = np.array([0,0,0])
        c, s = np.cos(anglebias), np.sin(anglebias)
        head_rot = np.array([
            [1,  0,   0],
            [0,  c,  -s],
            [0,  s,   c]
        ])
        hT_e[:3, 3] = -head_rot.T@head_pos
        hT_e[:3, :3] = head_rot.T


        rd_pos, rd_quat, ld_pos, ld_quat = pick_box_r_l(box_width, box_pos, box_quat)

        roT_e = roT_h@hT_e
        roP_D, roR_D = eD2oD(rd_pos, rd_quat, roT_e)

        loT_e = loT_h@hT_e
        loP_D, loR_D = eD2oD(ld_pos, ld_quat, loT_e)


        
        
        
        ctrl_l = motor_angle_clip(ik_general_L(q_real_float_l, loP_D, loR_D))
        ctrl_r = motor_angle_clip(ik_general(q_real_float_r, roP_D, roR_D))
        # ctrl_l = ctrl_r.copy()
        # ctrl_l[0] = -ctrl_r[0].copy()
        # ctrl_l[5] = np.pi - ctrl_r[5].copy()
        # ctrl_l[0] = -ctrl_r[0] #TODO 临时补救
        # ctrl_l = np.zeros(6)
        # ctrl_r = np.zeros(6)

        intikq_l = trans_q2intq(ctrl_l)
        intikq_r = trans_q2intq(ctrl_r)


        ctrl_intq_l = get_interpoint_l(q_real_int_l, intikq_l, max_dq)
        ctrl_intq_r = get_interpoint_r(q_real_int_r, intikq_r, max_dq)

        ctrl_intq = np.hstack((ctrl_intq_l, ctrl_intq_r))
        ctrl_q_degree = trans_intq2q_degree(ctrl_intq).tolist()
        # 添加条件限制打印频率
        if cnt % 300 == 0:  # 每300次循环打印一次
            print(f"Average time cost: {(time.time()-loopStartTime)/cnt:.4f}s")
            # print(f"q_real is {q_real_degree}")
            # print(f"ctrl_q_degree is {ctrl_q_degree}")
            print(f"ctrl l {ctrl_l} and ctrl_r {ctrl_r}")
            # print(f"site_force is {site_force}")
            print(f"width force bias = {clip_k*(clip_force - clip_force_0)*np.abs(clip_force - clip_force_0)}")
            print(f"box_width_0 is {box_width_0}")
            print(f"box_width is {box_width}")
            print(f"clip_force = {clip_force}")
            print(f"getbox_flag = {getbox_flag}")

        # 2.2 TODO:Compute U from Y
        ts = int(time.time() * 1_000_000)  # current timestamp from epoch (microsecond)
        u_data = UData(timestamp=ts, pos=ctrl_q_degree)

        # 2.3 Send U to remoteIP
        enet.SendU(u_data)

        # Sleep to maintain 1kHz loop
        elapsed_time = time.time() - start_time
        # print(elapsed_time)
        time.sleep(max(0, loopPeriod - elapsed_time))
except KeyboardInterrupt:
    print("Exiting...")